// implement fork from user space

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
static void
pgfault(struct UTrapframe *utf)
{
	void *addr = (void *) utf->utf_fault_va;
	uint32_t err = utf->utf_err;
	int r;

	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at uvpt
	//   (see <inc/memlayout.h>).

	// LAB 4: Your code here.
	pte_t pte;
	envid_t cur_envid;
	uint32_t index;
	uintptr_t aligned_addr;
	index = PDX(addr) * NPDENTRIES + PTX(addr);
	aligned_addr = ROUNDDOWN((uintptr_t)addr, PGSIZE);
	pte = uvpt[index];
	if (!(err & FEC_WR))
	{
		cprintf("pte is 0x%x addr 0x%x err %d\n", pte, addr, err);
		panic("pagefault is not caused by write\n");
	}
	if ((pte | PTE_COW) != pte)
	{
		cprintf("pte is 0x%x addr 0x%x\n", pte, addr);
		panic("pagefault not cow envid 0x%x\n", thisenv->env_id);
	}
	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.

	// LAB 4: Your code here.
	cur_envid = sys_getenvid();
	if (sys_page_alloc(cur_envid, PFTEMP, PTE_W | PTE_U | PTE_P))
	{
		panic("sys_page_alloc failed\n");
	}
	memcpy(PFTEMP, (void *)aligned_addr, PGSIZE);
	if (sys_page_map(cur_envid, PFTEMP, cur_envid, (void *)aligned_addr, PTE_W | PTE_U | PTE_P))
	{
		panic("sys_page_map failed in pagefault\n");
	}
	sys_page_unmap(cur_envid, PFTEMP);
	return;
	panic("pgfault not implemented");
}

//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why do we need to mark ours
// copy-on-write again if it was already copy-on-write at the beginning of
// this function?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
//
static int
duppage(envid_t envid, unsigned pn)
{
	int r;

	// LAB 4: Your code here.
	unsigned i = 0;
	pde_t pde;
	pte_t pte;
	envid_t cur_envid;
	uint32_t addr;
	cur_envid = sys_getenvid();
	for (i = 0; i < pn; ++i)
	{
		addr = i << PGSHIFT;
		if (addr == UXSTACKTOP - PGSIZE)
		{
			// 跳过异常处理栈复制，每个进程的异常处理栈独立分配
			continue;
		}
		pde = uvpd[PDX(addr)];
		if (pde & PTE_P)
		{
			pte = uvpt[PDX(addr) * NPDENTRIES + PTX(addr)];
			if ((pte & PTE_P) && (pte & PTE_U))
			{
				// cprintf("addr 0x%x index %d\n", addr, PDX(addr) * NPDENTRIES + PTX(addr));
				// cprintf("i %d start pde 0x%x pte 0x%x\n", i, pde, pte);
				if (((pte | PTE_W) == pte) || ((pte | PTE_COW) == pte))
				{
					// 当前pte具有写权限或者COW
					// 将当前进程addr处映射复制到子进程
					// cprintf("i %d cow or write\n", i);
					if (sys_page_map(cur_envid, (void *)addr, envid, (void *)addr, PTE_COW | PTE_U | PTE_P))
					{
						panic("sys_page_map cow failed in duppage child addr 0x%x pte 0x%x env-pgdir 0x%x thisenvid 0x%x pdx 0x%x ptx 0x%x real-ptx 0x%x\n", addr, pte, thisenv->env_pgdir, cur_envid, PDX(addr), PTX(addr), ((uintptr_t)addr >> 12) & (uintptr_t)0x3ff);
					}
					// cprintf("i %d sys_page_map cur_envid + envid finish\n", i);
					// 将当前进程addr处映射重新设置为COW
					if (sys_page_map(cur_envid, (void *)addr, cur_envid, (void *)addr, PTE_COW | PTE_U | PTE_P))
					{
						panic("sys_page_map cow failed in duppage current  addr 0x%x pte 0x%x\n", addr, pte);
					}
					// cprintf("i %d sys_page_map cur_envid finish\n", i);
				}
				else
				{
					// cprintf("i %d else\n", i);
					if (sys_page_map(cur_envid, (void *)addr, envid, (void *)addr, pte & PTE_SYSCALL))
					{
						cprintf("i is %d pte 0x%x\n", i, pte);
						panic("sys_page_map normal failed in duppage child\n");
					}
				}
				// cprintf("i %d finished pde 0x%x pte 0x%x\n", i, pde, pte);
			}
		}
	}
	// panic("duppage not implemented");
	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use uvpd, uvpt, and duppage.
//   Remember to fix "thisenv" in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
	// LAB 4: Your code here.
	envid_t cur_envid, child_envid;
	cur_envid = sys_getenvid();

	// cprintf("USTACKTOP 0x%x UXSTACKTOP 0x%x UTOP 0x%x\n", USTACKTOP, UXSTACKTOP, UTOP);
	// 1. The parent installs pgfault() as the C-level page fault handler, using the set_pgfault_handler() function you implemented above.
	set_pgfault_handler(pgfault);

	// 2. The parent calls sys_exofork() to create a child environment.
	child_envid = sys_exofork();
	if (child_envid < 0)
	{
		panic("fork failed at sys_exofork\n");
	}

	if (child_envid == 0)
	{
		thisenv = envs + ENVX(sys_getenvid());
		return 0;
	}

	// 3.For each writable or copy-on-write page in its address space below UTOP, the parent calls duppage, 
	// which should map the page copy-on-write into the address space of the child and then remap the page copy-on-write in its own address space.
	// [Note: The ordering here (i.e.,marking a page as COW in the child before marking it in the parent) actually matters! Can you see why? 
	// Try to think of a specific case where reversing the order could cause trouble. ]
	// duppage sets both PTEs so that the page is not writeable, and to contain PTE_COW in the "avail" field to distinguish 
	// copy-on-write pages from genuine read-only pages.
	duppage(child_envid, ((uintptr_t)UTOP) >> PGSHIFT);
	// The exception stack is not remapped this way, however. Instead you need to allocate a fresh page in the child for the exception stack. 
	// Since the page fault handler will be doing the actual copying and the page fault handler runs on the exception stack, 
	// the exception stack cannot be made copy-on-write: who would copy it?
	if (sys_page_alloc(child_envid, (void *)(UXSTACKTOP - PGSIZE), PTE_P | PTE_U | PTE_W))
	{
		panic("fork failed at exception stack alloc\n");
	}

	// 4. The parent sets the user page fault entrypoint for the child to look like its own.
	extern void _pgfault_upcall(void);
	sys_env_set_pgfault_upcall(child_envid, _pgfault_upcall);

	// 5. The child is now ready to run, so the parent marks it runnable.
	sys_env_set_status(child_envid, ENV_RUNNABLE);

	return child_envid;
	panic("fork not implemented");
}

// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}
